import socket
import cv2
import numpy as np
import threading
import time

class CameraClient:
    def __init__(self, host='192.168.4.1', port=8080):
        self.host = host
        self.port = port
        self.socket = None
        self.connected = False
        self.running = False
        
    def connect(self):
        """连接到ESP32摄像头服务器"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(10.0)
            self.socket.connect((self.host, self.port))
            self.connected = True
            print(f"Connected to {self.host}:{self.port}")
            return True
        except Exception as e:
            print(f"Connection failed: {e}")
            return False
    
    def disconnect(self):
        """断开连接"""
        self.running = False
        self.connected = False
        if self.socket:
            self.socket.close()
            self.socket = None
        print("Disconnected")
    
    def receive_metadata(self):
        """接收图像元数据"""
        try:
            meta_line = b""
            timeout = time.time() + 5.0
            while time.time() < timeout:
                byte = self.socket.recv(1)
                if byte == b'\n' or byte == b'':
                    break
                meta_line += byte
            
            if not meta_line:
                return None
                
            meta_str = meta_line.decode('utf-8').strip()
            # print(f"Raw metadata: {meta_str}") # 可以暂时注释掉减少输出
            if meta_str.startswith("META,"):
                parts = meta_str[5:].split(',')
                if len(parts) >= 3:
                    metadata = {
                        'row_size': int(parts[0]),
                        'width': int(parts[1]),
                        'height': int(parts[2])
                    }
                    # print(f"Parsed metadata: {metadata}")
                    return metadata
            return None
        except Exception as e:
            print(f"Error receiving metadata: {e}")
            return None
    
    def receive_image_data(self, row_size, height):
        """接收图像数据"""
        try:
            total_bytes = row_size * height
            # print(f"Expected bytes: {total_bytes}")
            
            data = b""
            while len(data) < total_bytes:
                chunk = self.socket.recv(min(4096, total_bytes - len(data)))
                if not chunk:
                    break
                data += chunk
                # print(f"Received: {len(data)}/{total_bytes} bytes", end='\r')
            
            # print(f"\nActual received bytes: {len(data)}")
            
            end_marker = b""
            timeout = time.time() + 2.0
            while len(end_marker) < 5 and time.time() < timeout:
                byte = self.socket.recv(1)
                if not byte:
                    break
                end_marker += byte
            
            # print(f"End marker: {end_marker}")
            
            if end_marker != b"\nEND\n":
                print("Warning: End marker not found or incorrect")
            
            return data if len(data) == total_bytes else None
        except Exception as e:
            print(f"Error receiving image data: {e}")
            return None
    
    def process_image(self, data, width, height):
        try:
            expected_size = width * height * 2
            actual_size = len(data)
            
            if actual_size != expected_size:
                print(f"Data size mismatch: got {actual_size}, expected {expected_size}")
                return None
                
            img_array = np.frombuffer(data, dtype=np.uint16)
            
            # 重新组织为 uint16 (小端序)
            img_array = img_array.byteswap()
            # RGB565 解码
            r = ((img_array >> 11) & 0x1F).astype(np.uint8) << 3 
            g =  ((img_array >> 5) & 0x3F).astype(np.uint8) << 2  
            b = (img_array & 0x1F).astype(np.uint8) << 3 
            # 创建RGB图像
            rgb_image = np.stack([r, g, b], axis=1).reshape((height, width, 3)).astype(np.uint8)
            
            # 转换为BGR用于OpenCV显示
            bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
            return bgr_image
            
        except Exception as e:
            print(f"Error processing image: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def start_receiving(self):
        """开始接收图像"""
        if not self.connected:
            print("Not connected to server")
            return
        
        self.running = True
        print("Starting to receive images...")
        
        frame_count = 0
        start_time = time.time()
        
        try:
            while self.running and self.connected:
                metadata = self.receive_metadata()
                if not metadata:
                    print("Failed to receive metadata")
                    break
                
                image_data = self.receive_image_data(
                    metadata['row_size'], 
                    metadata['height']
                )
                
                if not image_data:
                    print("Failed to receive image data")
                    break
                
                image = self.process_image(
                    image_data, 
                    metadata['width'], 
                    metadata['height']
                )
                
                if image is not None:
                    frame_count += 1
                    if frame_count % 30 == 0:
                        elapsed = time.time() - start_time
                        fps = frame_count / elapsed if elapsed > 0 else 0
                        print(f"FPS: {fps:.2f}")
                    
                    cv2.imshow('ESP32 Camera - Press Q to quit', image)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
                else:
                    print("Failed to process image")
        
        except KeyboardInterrupt:
            print("Interrupted by user")
        except Exception as e:
            print(f"Error during reception: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.disconnect()
            cv2.destroyAllWindows()
            print(f"Total frames: {frame_count}")

def main():
    client = CameraClient()
    
    if client.connect():
        try:
            client.start_receiving()
        except KeyboardInterrupt:
            print("Stopping client...")
        finally:
            client.disconnect()
    else:
        print("Failed to connect to camera server")

if __name__ == "__main__":
    main()